/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
 * Description: hccl cpu侧对外API接口实现文件
 * Author: gaozhenyeu
 * Create: 2024-04-02
 */

#include <atomic>
#include <algorithm>
#include <arpa/inet.h>
#include <fstream>
#include <fcntl.h>
#include <unistd.h>
#include <hccl/hccl_types.h>

#include "network/hccp.h"
#include "externalinput_pub.h"
#include "topoinfo_ranktableParser_pub.h"
#include "alg_remote_lookup.h"
#include "hdds_param.h"
#include "device_capacity.h"
#include "transport_heterog_roce_pub.h"
#include "hccl_impl_cpu.h"
#include "hccl_comm_cpu.h"

namespace hccl {

HcclCommCpu::HcclCommCpu(u64 inCCLbufferSize, u64 outCCLbufferSize, std::string identifier)
    : HcclCommBase(inCCLbufferSize, outCCLbufferSize, identifier)
{
}

HcclCommCpu::~HcclCommCpu()
{
    impl_ = nullptr;
}

HcclResult HcclCommCpu::CreateGroup(const std::string &group, const u32 &groupRank, const u32 &userRank,
    const std::vector<u32> &groupRanks, std::shared_ptr<HcclCommBase> &groupComm)
{
    // 增加输出日志关键字
    HCCL_INFO("HCCL_KEY_INFO: group[%s], groupRank[%u], userRank[%u], groupComm[%p]", group.c_str(),
        groupRank, userRank, groupComm.get());

    // 入参有消息校验
    if (group.length() == 0) {
        HCCL_ERROR("[Create][Group]errNo[0x%016llx] group name lenth is 0", HCCL_ERROR_CODE(HCCL_E_PARA));
        return HCCL_E_PARA;
    }

    if (groupRank >= groupRanks.size()) {
        HCCL_ERROR("[Create][Group]errNo[0x%016llx] group rank[%u] out of range [0,%llu])",
            HCCL_ERROR_CODE(HCCL_E_PARA), groupRank, groupRanks.size() - 1);
        return HCCL_E_PARA;
    }

    HcclRootInfo id;
    CHK_RET(GetUniqueId(&id));

    HcclCommParams params;
    params.rank = groupRank;
    params.userRank = userRank;
    params.totalRanks = groupRanks.size();
    params.isHeterogComm = isHeterogComm_;
    s32 iret = snprintf_s(params.id.internal, HCCL_ROOT_INFO_BYTES, HCCL_ROOT_INFO_BYTES - 1, "%s%s%s",
                          id.internal, "-", group.c_str());

    CHK_PRT_RET((iret == -1), HCCL_ERROR("[Create][Group]errNo[0x%016llx] get group unique id falied",
        HCCL_ERROR_CODE(HCCL_E_INTERNAL)), HCCL_E_INTERNAL);

    WorldGroupInfo groupCommonData;

    CHK_RET(impl_->GetGroupCommonData(groupCommonData));
    params.logicDevId = groupCommonData.deviceLogicId;
    params.profilingInitiated = groupCommonData.profilingInitiated;
    params.deviceType = deviceType_;
    params.hcomGroupNicInit = impl_->GetNicInitialized();
    std::vector<RankInfo> rankList;

    CHK_RET(impl_->GetGroupRanksInfo(groupRanks, rankList));

    groupComm.reset(new (std::nothrow) hccl::HcclCommCpu(0, 0, group));
    CHK_SMART_PTR_NULL(groupComm);
    CHK_RET(groupComm->init(params, rankList, groupCommonData));

    return HCCL_SUCCESS;
}

HcclResult HcclCommCpu::InitImpl(DevType deviceType)
{
    impl_.reset((new (std::nothrow) HcclImplCpu()));
    deviceType_ = deviceType;
    return HCCL_SUCCESS;
}

HcclResult HcclCommCpu::send(const std::string &tag, void *inputPtr, u64 count, HcclDataType dataType,
    u32 destRank, rtStream_t stream)
{
    /* 增加输出日志关键字 */
    HCCL_DEBUG("HCCL_KEY_INFO: tag[%s], input_ptr[%p], count[%llu], data_type[%s], destRank[%u]",
        tag.c_str(), inputPtr, count, GetDataTypeEnumStr(dataType).c_str(), destRank);

    /* 入参检查 */
    CHK_PTR_NULL(inputPtr);

    if (tag.empty()) {
        HCCL_ERROR("[HcclComm][Send]errNo[0x%016llx] send tag length is 0",
            HCCL_ERROR_CODE(HCCL_E_PARA));
        return HCCL_E_PARA;
    }

    CHK_RET(impl_->CheckCount(count));
    CHK_RET(impl_->CheckDataType(dataType, false));
    CHK_RET(impl_->CheckUserRank(destRank));
    CHK_RET(reinterpret_cast<hccl::HcclImplCpu *>(impl_.get())->Send(tag, inputPtr, count, dataType, destRank, stream));

    return HCCL_SUCCESS;
}

HcclResult HcclCommCpu::receive(const std::string &tag, void *outputPtr, u64 count, HcclDataType dataType,
    u32 srcRank, rtStream_t stream)
{
    /* 增加输出日志关键字 */
    HCCL_DEBUG("HCCL_KEY_INFO: tag[%s], output_ptr[%p], count[%llu], data_type[%s], srcRank[%u]",
               tag.c_str(), outputPtr, count, GetDataTypeEnumStr(dataType).c_str(), srcRank);

    /* * 入参检查 */
    CHK_PTR_NULL(outputPtr);

    CHK_PRT_RET(tag.empty(), HCCL_ERROR("[HcclComm][Receive]errNo[0x%016llx] receive tag length is 0",
        HCCL_ERROR_CODE(HCCL_E_PARA)), HCCL_E_PARA);

    CHK_RET(impl_->CheckCount(count));
    CHK_RET(impl_->CheckDataType(dataType, false));
    CHK_RET(impl_->CheckUserRank(srcRank));
    CHK_RET(reinterpret_cast<hccl::HcclImplCpu *>(impl_.get())->
        Receive(tag, outputPtr, count, dataType, srcRank, stream));

    return HCCL_SUCCESS;
}

HcclResult HcclCommCpu::PrepareStartColl(const std::string &tag, const HcomCollOpInfo &opInfo, HcomOperationType opType)
{
    /* 入参检查 */
    CHK_PRT_RET(tag.empty(), HCCL_ERROR("[HcclComm][PrepareStartColl]op tag length is 0"), HCCL_E_PARA);

    CHK_RET(impl_->CheckCount(opInfo.count));
    CHK_RET(impl_->CheckDataType(opInfo.dataType, false));

    CHK_RET(reinterpret_cast<hccl::HcclImplCpu *>(impl_.get())->PrepareStartColl(tag, opInfo, opType));
    return HCCL_SUCCESS;
}

HcclResult HcclCommCpu::PrepareStart(const std::string &tag, void *addr, u64 count, HcclDataType dataType,
    u32 destRank)
{
    /* 入参检查 */
    CHK_PRT_RET(tag.empty(), HCCL_ERROR("[HcclComm][PrepareStart]op tag length is 0"), HCCL_E_PARA);

    CHK_RET(impl_->CheckCount(count));
    CHK_RET(impl_->CheckDataType(dataType, false));
    CHK_RET(impl_->CheckUserRank(destRank));
    CHK_RET(reinterpret_cast<hccl::HcclImplCpu *>(impl_.get())->PrepareStart(tag, addr, count, dataType, destRank));
    return HCCL_SUCCESS;
}

HcclResult HcclCommCpu::PrepareQuerry(const std::string &tag, u32& status)
{
    /* 入参检查 */
    CHK_PRT_RET(tag.empty(), HCCL_ERROR("[HcclComm][PrepareStart]op tag length is 0"), HCCL_E_PARA);

    CHK_RET(reinterpret_cast<hccl::HcclImplCpu *>(impl_.get())->PrepareQuerry(tag, status));
    return HCCL_SUCCESS;
}

HcclResult HcclCommCpu::PrepareCancel(const std::string &tag, u32& status)
{
    /* 入参检查 */
    CHK_PRT_RET(tag.empty(), HCCL_ERROR("[HcclComm][PrepareCancel]op tag length is 0"), HCCL_E_PARA);

    CHK_RET(reinterpret_cast<hccl::HcclImplCpu *>(impl_.get())->PrepareCancel(tag, status));
    return HCCL_SUCCESS;
}


HcclResult HcclCommCpu::GatherByOs(void* inputBuf, uint64_t inputCount, HcclDataType inputType, void* outputBuf,
    uint64_t outputCount, HcclDataType outputType, int root, uint64_t flag, const std::string &tag)
{
        /* 增加输出日志关键字 */
    HCCL_INFO("HCCL_KEY_INFO: input_ptr[%p], output_ptr[%p], count[%llu], data_type[%s]", inputBuf,
        outputBuf, inputCount, GetDataTypeEnumStr(inputType).c_str());

    /* * 入参检查 */
    CHK_PTR_NULL(inputBuf);
    CHK_PRT_RET(tag.empty(), HCCL_ERROR("[HcclComm][Gather]errNo[0x%016llx] gather tag length is 0",
        HCCL_ERROR_CODE(HCCL_E_PARA)), HCCL_E_PARA);

    CHK_RET(impl_->CheckCount(inputCount));
    CHK_RET(impl_->CheckDataType(inputType, false));

    CHK_RET(reinterpret_cast<hccl::HcclImplCpu *>(impl_.get())->
        GatherByOs(inputBuf, inputCount, inputType, outputBuf, outputCount, outputType, root, flag, tag));
    
    return HCCL_SUCCESS;
}

HcclResult HcclCommCpu::BcastByOS(void* buf, uint64_t count, HcclDataType dataType,
    int root, uint64_t flag, const std::string &tag)
{
    /* * 入参检查 */
    CHK_PTR_NULL(buf);
    CHK_PRT_RET(tag.empty(), HCCL_ERROR("[HcclComm][BcastByOS]errNo[0x%016llx] BcastByOS tag length is 0",
        HCCL_ERROR_CODE(HCCL_E_PARA)), HCCL_E_PARA);

    /* 增加输出日志关键字 */
    HCCL_INFO("HCCL_KEY_INFO: buf[%p], count[%llu], data_type[%s]", buf, count, GetDataTypeEnumStr(dataType).c_str());

    CHK_RET(impl_->CheckCount(count));
    CHK_RET(impl_->CheckDataType(dataType, false));
    CHK_RET(impl_->CheckUserRank(root));

    CHK_RET(reinterpret_cast<hccl::HcclImplCpu *>(impl_.get())->BcastByOS(buf, count, dataType, root, flag, tag));

    return HCCL_SUCCESS;
}

HcclResult HcclCommCpu::GetHeterogCollTag(HcomOperationType opType, std::string &tag)
{
    CHK_RET(reinterpret_cast<hccl::HcclImplCpu *>(impl_.get())->GetHeterogCollTag(opType, tag));
    return HCCL_SUCCESS;
}

}